# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: broker\worker_manager\workers_storage\__init__.py
__author__ = "Attila Gerendi (Sanyi)"
import logging, datetime
from collections import defaultdict
from typing import Union
from sqlalchemy import select as sql_select
from db import Connection
from db.tables.workers import WorkerRow, WorkersTable, WorkerTypes, WorkerStatus
from helpers.http_clients.scan_api import ScanApiClient
from scanners.scan_app.errors import WorkerNotFound, JobNotFound, WorkerNotAuthorized
from db.tables.workers import WorkerAuthorization
from broker.worker_manager.jobs_storage import JobsStorage

_logger = logging.getLogger("broker.worker_manager.workers_storage")


class WorkersStorage:
    def __init__(self, shard):
        super().__init__()
        self.shard = shard

    def get_scan_api_client(self, worker_id, request_timeout=20):
        worker = self.get(worker_id)
        if not worker:
            raise WorkerNotFound()
        if worker["authorization"] != WorkerAuthorization.AUTHORIZED:
            raise WorkerNotAuthorized()
        client = ScanApiClient(
            endpoint=worker["endpoint"],
            scanning_app=worker["scanning_app"],
            api_key=worker["api_key"],
            ssl_verify=worker["ssl_verify"],
            request_timeout=request_timeout,
        )
        return client

    def get_scan_api_client_by_job(self, jobs_storage, scan_session_id, scanning_app):
        job = jobs_storage.get(scan_session_id, scanning_app)
        if not job:
            _logger.error(
                "no descriptor found for %s:%s", scan_session_id, scanning_app
            )
            raise JobNotFound()
        if job["worker_id"] == "ffffffff-ffff-ffff-ffff-ffffffffffff":
            return
        else:
            return self.get_scan_api_client(job["worker_id"])

    def get(self, worker_id):
        q = sql_select((WorkerRow,)).where(WorkerRow.worker_id == worker_id)
        with Connection(self.shard) as (db):
            worker = db.execute(q).fetchone()
            if not worker:
                raise WorkerNotFound()
            return dict(worker)

    def get_worker_ids_by_scanning_app(self, scanning_app, expected_status):
        q = sql_select((WorkerRow.worker_id,)).where(
            WorkerRow.scanning_app == scanning_app
        )
        if isinstance(expected_status, str):
            q = q.where(WorkerRow.status == expected_status)
        else:
            if isinstance(expected_status, (list, tuple)):
                for status in expected_status:
                    q = q.where(WorkerRow.status == status)

        with Connection(self.shard) as (db):
            result = []
            for r in db.execute(q).fetchall():
                result.append(r.worker_id)

            return result

    def get_all_workers(self):
        q = sql_select((WorkerRow,))
        workers = dict()
        with Connection(self.shard) as (db):
            for worker in db.execute(q).fetchall():
                workers[worker.worker_id] = dict(worker)

        return workers

    def get_all_auto_workers(self):
        q = sql_select((WorkerRow,)).where(WorkerRow.worker_type == "auto")
        workers = []
        with Connection(self.shard) as (db):
            for worker in db.execute(q).fetchall():
                workers.append(dict(worker))

        return workers

    def get_workers_by_scanning_app_and_status(self):
        workers = defaultdict(lambda: defaultdict(list))
        q = sql_select((WorkerRow,))
        with Connection(self.shard) as (db):
            for worker in db.execute(q).fetchall():
                workers[worker.scanning_app][worker.status].append(dict(worker))

        return workers

    def set_worker_stating(self, worker_id):
        q = (
            WorkersTable.update(
                values=dict(
                    status=WorkerStatus.STARTING, start_date=datetime.datetime.utcnow()
                )
            )
        ).where(WorkerRow.worker_id == worker_id)
        with Connection(self.shard) as (db):
            db.execute(q)

    def set_worker_stopping(self, worker_id):
        q = (
            WorkersTable.update(
                values=dict(
                    status=WorkerStatus.STOPPING,
                    shutdown_date=datetime.datetime.utcnow(),
                )
            )
        ).where(WorkerRow.worker_id == worker_id)
        with Connection(self.shard) as (db):
            db.execute(q)

    def update_workers_status(self, worker_id_list, status, only_auto_workers=True):
        q = (WorkersTable.update(values=dict(status=status))).where(
            WorkerRow.worker_id.in_(worker_id_list)
        )
        if only_auto_workers:
            q = q.where(WorkerRow.worker_type == WorkerTypes.AUTO)
        with Connection(self.shard) as (db):
            db.execute(q)
